package com.dycong.common.http.SSLEngineUtil;

import org.apache.log4j.Logger;

import javax.annotation.Resource;
import javax.net.ssl.*;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateException;

/**
 * Created by dycong on 2017/4/3.
 */
public abstract class SSLEnginePeer {

    protected final Logger log = Logger.getLogger(SSLEnginePeer.class);
    @Resource
    SSLExecutor sslExecutor=new SSLExecutor();

    protected ByteBuffer inAppBuffer;
    protected ByteBuffer outAppBuffer;
    protected ByteBuffer inNetBuffer;
    protected ByteBuffer outNetBuffer;
    protected boolean bufferInit = false;

    public SSLEnginePeer() {

    }

    public SSLEnginePeer(int inAppSize, int outAppSize, int inNetSize, int outNetSize) {
        inAppBuffer = ByteBuffer.allocate(inAppSize);
        outAppBuffer = ByteBuffer.allocate(outAppSize);
        inNetBuffer = ByteBuffer.allocate(inNetSize);
        outNetBuffer = ByteBuffer.allocate(outNetSize);
        bufferInit = true;
    }

    protected abstract void read(SocketChannel socketChannel, SSLEngine sslEngine);

    protected void write(SocketChannel socketChannel, SSLEngine sslEngine, String message) {
        ByteBuffer b = ByteBuffer.allocate(message.length());
        b.put(message.getBytes());
        write(socketChannel, sslEngine, b);
    }

    protected abstract void write(SocketChannel socketChannel, SSLEngine sslEngine, ByteBuffer message);

    protected boolean handShake(SocketChannel socketChannel, SSLEngine sslEngine) {
        log.debug("handShake start...");
        SSLEngineResult result;
        SSLEngineResult.HandshakeStatus handshakeStatus;
        if (bufferInit) {
            inAppBuffer.clear();
            inNetBuffer.clear();
            outAppBuffer.clear();
            outNetBuffer.clear();
        } else {
            inAppBuffer = ByteBuffer.allocate(sslEngine.getSession().getApplicationBufferSize());
            outAppBuffer = ByteBuffer.allocate(sslEngine.getSession().getApplicationBufferSize());
            inNetBuffer = ByteBuffer.allocate(sslEngine.getSession().getPacketBufferSize());
            outNetBuffer = ByteBuffer.allocate(sslEngine.getSession().getPacketBufferSize());
            bufferInit = true;
        }
        handshakeStatus = sslEngine.getHandshakeStatus();
        while (handshakeStatus != SSLEngineResult.HandshakeStatus.FINISHED && handshakeStatus != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING) {
            switch (handshakeStatus) {
                case NEED_UNWRAP:
                    try {
                        socketChannel.read(inNetBuffer);
                        inNetBuffer.flip();
                        if (!doUnWrap(sslEngine))
                            return false;
                    } catch (IOException e) {
                        log.error("socketChannel read error..."+e);
                    }
                    handshakeStatus = sslEngine.getHandshakeStatus();
                    break;
                case NEED_WRAP:
                    outNetBuffer.clear();
                    if (!doWrap(sslEngine, socketChannel))
                        return false;
                    handshakeStatus = sslEngine.getHandshakeStatus();
                    break;
                case NEED_TASK:
                    doTask(sslEngine);
                    handshakeStatus = sslEngine.getHandshakeStatus();
                    break;
                case FINISHED:
                    break;
                case NOT_HANDSHAKING:
                    break;
                default:
                    throw new IllegalStateException("Invalid SSL status: " + handshakeStatus);

            }
        }
        return true;
    }

    protected boolean doTask(SSLEngine sslEngine) {
        boolean suc = true;
        Runnable runnable;
        try {
            while ((runnable = sslEngine.getDelegatedTask()) != null) {
                sslExecutor.execute(runnable);
                log.info("task run...");
            }
        } catch (Exception e) {
            suc = false;
            log.error("task run error" + e);
        }
        return suc;
    }

    protected boolean doUnWrap(SSLEngine sslEngine) {
        boolean suc = true;
        try {
                SSLEngineResult result = sslEngine.unwrap(inNetBuffer, inAppBuffer);
            // TODO: 2017/4/4 如果此时 UNDERFLOW，在compact之后在进行handleBufferUnderflow()方法-->有选择的扩容
                inNetBuffer.compact();

                switch (result.getStatus()) {
                    case OK:
                        break;
                    case BUFFER_OVERFLOW:
                        inAppBuffer = enlargeApplicationBuffer(sslEngine, inAppBuffer);
                        break;
                    case BUFFER_UNDERFLOW:
                        inNetBuffer = handleBufferUnderflow(sslEngine, inNetBuffer);
                        break;
                    case CLOSED:
                        if (sslEngine.isOutboundDone()) {
                            return false;
                        } else {
                            sslEngine.closeOutbound();
                            break;
                        }
                    default:
                        throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
                }
            //todo 这里unwrap循环的话，无法socket，read()自然死循环 } while (status == SSLEngineResult.HandshakeStatus.NEED_UNWRAP);
            log.info("sslEngine unwrap...");
        } catch (SSLException e) {
            suc = false;
            sslEngine.closeOutbound();
            log.error("unwrap error：" + e);
        }
        return suc;
    }

    protected boolean doWrap(SSLEngine sslEngine, SocketChannel socketChannel) {
        boolean suc = true;
        SSLEngineResult.HandshakeStatus status = null;
        try {
                SSLEngineResult result = sslEngine.wrap(outAppBuffer, outNetBuffer);

                switch (result.getStatus()) {
                    case OK:
                        try {
                            outNetBuffer.flip();
                            while (outNetBuffer.hasRemaining()) {
                                socketChannel.write(outNetBuffer);
                            }
                        } catch (IOException e) {
                            log.error("write error: " + e);
                        }

                        break;
                    case BUFFER_OVERFLOW:
                        outAppBuffer = enlargePacketBuffer(sslEngine, outNetBuffer);
                        break;
                    case BUFFER_UNDERFLOW:
                        throw new IllegalArgumentException("Buffer underflow occured after a wrap");
                    case CLOSED:
                        try {
                            outNetBuffer.flip();
                            while (outNetBuffer.hasRemaining()) {
                                socketChannel.write(outNetBuffer);
                            }
                            //todo At this point the handshake status will probably be NEED_UNWRAP so we make sure that peerNetData is clear to read.
                            inNetBuffer.clear();
                        } catch (IOException e) {
                            log.error("Failed to send server's CLOSE message due to socket channel's failure.");
                        }
                        break;
                    default:
                        throw new IllegalStateException("Invalid SSL status: " + result.getStatus());
                }
                log.info("sslEngine wrap...");
        } catch (SSLException e) {
            suc = false;
            sslEngine.closeOutbound();
            log.info("wrap error..." + e);
        }
        return suc;
    }

    protected ByteBuffer enlargePacketBuffer(SSLEngine engine, ByteBuffer buffer) {
        return enlargeBuffer(buffer, engine.getSession().getPacketBufferSize());
    }

    protected ByteBuffer enlargeApplicationBuffer(SSLEngine engine, ByteBuffer buffer) {
        return enlargeBuffer(buffer, engine.getSession().getApplicationBufferSize());
    }

    protected ByteBuffer enlargeBuffer(ByteBuffer buffer, int sessionProposedCapacity) {
        if (sessionProposedCapacity > buffer.capacity()) {
            buffer = ByteBuffer.allocate(sessionProposedCapacity);
        } else {
            buffer = ByteBuffer.allocate(buffer.capacity() * 2 + 1);
        }
        return buffer;
    }

    //当inNetPeer一次数据大于最初设定的cap时，不能简单的“*2”因为在stream只能读一次，要保留数据以便再继续读
    protected ByteBuffer handleBufferUnderflow(SSLEngine sslEngine, ByteBuffer buffer) {
        if (sslEngine.getSession().getPacketBufferSize() < buffer.limit()) {
            return buffer;
        } else {
            ByteBuffer replaceBuffer = enlargePacketBuffer(sslEngine, buffer);
            buffer.flip();
            replaceBuffer.put(buffer);
            return replaceBuffer;
        }
    }

    protected void closeConnection(SocketChannel socketChannel, SSLEngine sslEngine) {
        sslEngine.closeOutbound();
        handShake(socketChannel, sslEngine);
        try {
            socketChannel.close();
        } catch (IOException e) {
            log.error("socket outBound close error: " + e);
        }
    }

    //当对方关闭输出流时，我方可以关闭输入流，read()返回值为-1时
    protected void handleEndOfStream(SocketChannel socketChannel, SSLEngine sslEngine) {
        try {
            sslEngine.closeInbound();
        } catch (IOException e) {
            log.error("socket inbound close error" + e);
        }
        //todo 关闭断开的socketChannel，以及废弃的、只能用一次的sslEngine
        closeConnection(socketChannel, sslEngine);
    }

    protected KeyManager[] createKeyManagers(String filepath, String keyStoryPassWord, String keyPassWord) {
        KeyManagerFactory kmf = null;
        try (FileInputStream in = new FileInputStream(filepath)) {
            KeyStore keyStore = KeyStore.getInstance("JKS");
            keyStore.load(in, keyStoryPassWord.toCharArray());
            kmf = KeyManagerFactory.getInstance("SunX509");
            kmf.init(keyStore, keyPassWord.toCharArray());
        } catch (KeyStoreException e) {
            log.error("keyStory instance error: " + e);
        } catch (FileNotFoundException e) {
            log.error("KeyStore file no find: " + e);
        } catch (IOException e) {
            log.error("" + e);
        } catch (NoSuchAlgorithmException e) {
            log.error("" + e);
        } catch (CertificateException e) {
            log.error("" + e);
        } catch (UnrecoverableKeyException e) {
            log.error("kmf init error:" + e);
        }
        return kmf.getKeyManagers();
    }

    protected TrustManager[] createTrustManagers(String filepath, String keyStoryPassword) {
        TrustManagerFactory tmf = null;
        try (FileInputStream in = new FileInputStream(filepath)) {
            KeyStore trustStory = KeyStore.getInstance("JKS");
            trustStory.load(in, keyStoryPassword.toCharArray());
            tmf = TrustManagerFactory.getInstance("SunX509");
            tmf.init(trustStory);
        } catch (KeyStoreException e) {
            log.error("trustStory instance error: " + e);
        } catch (IOException e) {
            log.error("" + e);
        } catch (NoSuchAlgorithmException e) {
            log.error("" + e);
        } catch (CertificateException e) {
            log.error("" + e);
        }
        return tmf.getTrustManagers();
    }

}
